{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "在前置环节中，已经根据关键词从很多企业信息中筛选出包含“氢”作为关键词的企业"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "接下来依据企业信息使用大模型对企业细分到小类"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from transformers import AutoTokenizer\n",
    "from vllm import LLM, SamplingParams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-06-28 16:09:42,488 - modelscope - INFO - PyTorch version 2.3.0 Found.\n",
      "2024-06-28 16:09:42,489 - modelscope - INFO - Loading ast index from /home/jie/.cache/modelscope/ast_indexer\n",
      "2024-06-28 16:09:42,504 - modelscope - INFO - Loading done! Current index file version is 1.14.0, with md5 9f53011443b7fabdfb807148bea6e1f2 and a total number of 976 components indexed\n"
     ]
    }
   ],
   "source": [
    "from modelscope import snapshot_download"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = \"ZhipuAI/glm-4-9b-chat\"\n",
    "model_name = snapshot_download(model_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "\n",
    "\n",
    "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# os.environ['HTTP_PROXY'] = 'http://127.0.0.1:7890'\n",
    "# os.environ['HTTPS_PROXY'] = 'http://127.0.0.1:7890'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# filename = \"../data_clean/氢能_企业信息.csv\"\n",
    "# df = pd.read_csv(filename)\n",
    "# df.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from vllm import LLM, SamplingParams"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_from_disk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_dataset = load_from_disk(\"hydrogen_binary_dataset\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "60780"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(test_dataset[\"train\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(60780, 7)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_dataset[\"train\"].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['企业名称', '统一社会信用代码', '经营范围', '一级行业分类', '二级行业分类', '三级行业分类', 'labels'],\n",
       "    num_rows: 60780\n",
       "})"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_dataset[\"train\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# GLM-4-9B-Chat-1M\n",
    "# max_model_len, tp_size = 1048576, 4\n",
    "\n",
    "# GLM-4-9B-Chat\n",
    "# 如果遇见 OOM 现象，建议减少max_model_len，或者增加tp_size\n",
    "max_model_len, tp_size = 131072, 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = [\n",
    "            [{\"role\": \"user\", \"content\": test_dataset[\"train\"][i][\"labels\"]}]\n",
    "            for i in range(10)\n",
    "        ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[[{'role': 'user',\n",
       "   'content': '请你根据下述企业信息，判断该企业是否属于氢能企业：\\n    企业名称: 马鞍山市衡安化工产品贸易有限责任公司\\n\\n经营范围: 批发（不得存储）甲苯、重苯（苯渣）、焦化苯（轻苯）、二甲苯、洗油、溶剂油、甲醇、苯乙烯、煤焦油、萘、硫磺、蒽油、初馏份、脱酚酚油、甲酚（工业甲酚）、苯酚、二甲酚、邻甲酚、蒽、煤焦沥青、顺丁烯二酸酐、氢氧化钠、硫酸、粗苯、二甲残油、轻油；批发零售化工产品（除危险品、易制毒品）、水处理药剂、润滑油、工矿配件、机电产品、五金建材、消防器材、劳保用品、电子产品、生铁、耐火材料、钢材、办公用品。（依法须经批准的项目，经相关部门批准后方可开展经营活动）\\n\\n一级行业分类: 制造业\\n\\n二级行业分类: 石油、煤炭及其他燃料加工业\\n\\n三级行业分类: 煤炭加工\\n\\n\\n    请以list的形式返回结果，若是氢能企业，请返回[\"是\"]，否则返回[\"否\"]。\\n    '}],\n",
       " [{'role': 'user',\n",
       "   'content': '请你根据下述企业信息，判断该企业是否属于氢能企业：\\n    企业名称: 马鞍山市沛达商贸有限公司\\n\\n经营范围: 批发（不得储存）氢氧化钠、亚硝酸钠、磷酸、氢氟酸、双氧水、二甲苯、甲醛、甲醇、冰醋酸，批发及零售化工原料（不含危险化学品及易制毒品）、劳保用品、灯具、建材、金属材料、五金交电。（依法需经批准的项目经相关部门批准后方可经营）\\n\\n一级行业分类: 制造业\\n\\n二级行业分类: 化学原料和化学制品制造业\\n\\n三级行业分类: 基础化学原料制造\\n\\n\\n    请以list的形式返回结果，若是氢能企业，请返回[\"是\"]，否则返回[\"否\"]。\\n    '}],\n",
       " [{'role': 'user',\n",
       "   'content': '请你根据下述企业信息，判断该企业是否属于氢能企业：\\n    企业名称: 马鞍山市昌和染料有限公司\\n\\n经营范围: 批发零售染料、颜料、建材、五金工具、家用电器、机电产品、劳保用品、橡胶制品、化工原料（不含危险化学品及易制毒品）；零售（不得超量储存）、批发氢氧化钠、氢氟酸、亚硝酸钠、氟化氢铵、双氧水（含量22%）、氨水、漂粉精、漂白粉、甲醇、磷酸、氟硅酸钠、环氧树脂、无水乙醇、冰醋酸、过硫酸铵、硫脲、亚硫酸氢钠、氯化锌、异丙醇、甲醛、松香水、氨基磺酸、硫化碱。\\n\\n一级行业分类: 制造业\\n\\n二级行业分类: 化学原料和化学制品制造业\\n\\n三级行业分类: 基础化学原料制造\\n\\n\\n    请以list的形式返回结果，若是氢能企业，请返回[\"是\"]，否则返回[\"否\"]。\\n    '}],\n",
       " [{'role': 'user',\n",
       "   'content': '请你根据下述企业信息，判断该企业是否属于氢能企业：\\n    企业名称: 马鞍山市春林化工有限公司\\n\\n经营范围: 批发（不得储存、不得零售）氨水、氢氧化钠、盐酸、液氨；批发金属制品、建材、钢材、五金机械、电工电料、劳保用品 、百货日杂、装潢材料、润滑油、仪器仪表、塑料制品。（依法需经批准的项目经相关部门批准后方可开展经营活动）\\n\\n一级行业分类: 制造业\\n\\n二级行业分类: 化学原料和化学制品制造业\\n\\n三级行业分类: 基础化学原料制造\\n\\n\\n    请以list的形式返回结果，若是氢能企业，请返回[\"是\"]，否则返回[\"否\"]。\\n    '}],\n",
       " [{'role': 'user',\n",
       "   'content': '请你根据下述企业信息，判断该企业是否属于氢能企业：\\n    企业名称: 马鞍山克众贸易有限公司\\n\\n经营范围: 批发（不得储存）叔丁基过氧化氢、过氧化二叔丁基、过氧化苯甲酸叔丁酯、1,1-双（叔丁基过氧基）环已烷、1,1-双（叔丁基过氧基）-3,3,5-三甲基环已烷、苯甲酰氯、叔丁醇、氢氧化钠、硫酸、环已酮、过氧化二苯甲酰，批发建材、金属材料、五金交电、劳保用品、机械设备、电子产品、办公用品、日用百货。（依法须经批准的项目，经相关部门批准后方可开展经营活动）\\n\\n一级行业分类: 制造业\\n\\n二级行业分类: 化学原料和化学制品制造业\\n\\n三级行业分类: 基础化学原料制造\\n\\n\\n    请以list的形式返回结果，若是氢能企业，请返回[\"是\"]，否则返回[\"否\"]。\\n    '}],\n",
       " [{'role': 'user',\n",
       "   'content': '请你根据下述企业信息，判断该企业是否属于氢能企业：\\n    企业名称: 马鞍山市亚大化玻有限责任公司\\n\\n经营范围: 零售（不得超量储存）、批发：无水乙醇、95%乙醇、甲醇、氨水、氢氧化钠、二甲苯、四氯化碳、盐酸、硫酸、甲苯、高锰酸钾、丙酮、锌粒、氢氧化钾、化工产品（不含危险化学品及易制毒品）；玻璃仪器、五金交化、机电建材、橡胶制品、办公用品批发零售。（依法须经批准的项目，经相关部门批准后方可开展经营活动）\\n\\n一级行业分类: 制造业\\n\\n二级行业分类: 化学原料和化学制品制造业\\n\\n三级行业分类: 基础化学原料制造\\n\\n\\n    请以list的形式返回结果，若是氢能企业，请返回[\"是\"]，否则返回[\"否\"]。\\n    '}],\n",
       " [{'role': 'user',\n",
       "   'content': '请你根据下述企业信息，判断该企业是否属于氢能企业：\\n    企业名称: 马鞍山意尔佳贸易有限公司\\n\\n经营范围: 批发（不得储存）叔丁基过氧化氢、过氧化二叔丁基、过氧化苯甲酸叔丁酯、1，1-双（叔丁基过氧基）环已烷、1，1-双（叔丁基过氧基）-3，3，5-三甲基环已烷、苯甲酰氯、叔丁醇、氢氧化钠、硫酸、环已酮、过氧化二苯甲酰，批发建材、金属材料、五金交电、劳保用品、机械设备、电子产品、办公用品、日用百货。（依法须经批准的项目，经相关部门批准后方可开展经营活动）\\n\\n一级行业分类: 制造业\\n\\n二级行业分类: 化学原料和化学制品制造业\\n\\n三级行业分类: 基础化学原料制造\\n\\n\\n    请以list的形式返回结果，若是氢能企业，请返回[\"是\"]，否则返回[\"否\"]。\\n    '}],\n",
       " [{'role': 'user',\n",
       "   'content': '请你根据下述企业信息，判断该企业是否属于氢能企业：\\n    企业名称: 马鞍山瑞尔德贸易有限公司\\n\\n经营范围: 批发（不得储存）叔丁基过氧化氢、过氧化二叔丁基、过氧化苯甲酸叔丁酯、1，1-双（叔丁基过氧基）环已烷、1，1-双（叔丁基过氧基）-3，3，5-三甲基环已烷、苯甲酰氯、叔丁醇、氢氧化钠、硫酸、环已酮、过氧化二苯甲酰，批发建材、金属材料、五金交电、劳保用品、机械设备、电子产品、办公用品、日用百货。（依法须经批准的项目，经相关部门批准后方可开展经营活动）\\n\\n一级行业分类: 制造业\\n\\n二级行业分类: 化学原料和化学制品制造业\\n\\n三级行业分类: 基础化学原料制造\\n\\n\\n    请以list的形式返回结果，若是氢能企业，请返回[\"是\"]，否则返回[\"否\"]。\\n    '}],\n",
       " [{'role': 'user',\n",
       "   'content': '请你根据下述企业信息，判断该企业是否属于氢能企业：\\n    企业名称: 马鞍山市天源化工有限公司\\n\\n经营范围: 批发化工产品及原料（不含危险化学品及易制毒品）、建筑材料、矿石、五金交电、电线电缆、日用百货、机电设备、服装、鞋帽、防水防火材料、钢材,；批发（不得储存、不得零售）硫酸、烧碱、盐酸、精萘、煤焦油、苯、甲苯、二甲苯、甲醛、苯胺、亚硝酸钠、氢氧化钾、苯酐、液氨、二氧化硫、三氧化硫、氨水、氨基磺酸、甲醇、乙醇、氢氟酸、油漆。（依法须经批准的项目，经相关部门批准后方可开展经营活动）\\n\\n一级行业分类: 制造业\\n\\n二级行业分类: 化学原料和化学制品制造业\\n\\n三级行业分类: 基础化学原料制造\\n\\n\\n    请以list的形式返回结果，若是氢能企业，请返回[\"是\"]，否则返回[\"否\"]。\\n    '}],\n",
       " [{'role': 'user',\n",
       "   'content': '请你根据下述企业信息，判断该企业是否属于氢能企业：\\n    企业名称: 马鞍山沁丰贸易有限公司\\n\\n经营范围: 批发（不得储存）叔丁基过氧化氢、过氧化二叔丁基、过氧化苯甲酸叔丁酯、1，1-双（叔丁基过氧基）环已烷、1，1-双（叔丁基过氧基）-3，3，5-三甲基环已烷、苯甲酰氯、叔丁醇、氢氧化钠、硫酸、环已酮、过氧化二苯甲酰，批发建材、金属材料、五金交电、劳保用品、机械设备、电子产品、办公用品、日用百货。（依法须经批准的项目，经相关部门批准后方可开展经营活动）\\n\\n一级行业分类: 制造业\\n\\n二级行业分类: 化学原料和化学制品制造业\\n\\n三级行业分类: 基础化学原料制造\\n\\n\\n    请以list的形式返回结果，若是氢能企业，请返回[\"是\"]，否则返回[\"否\"]。\\n    '}]]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO 06-28 16:15:26 llm_engine.py:161] Initializing an LLM engine (v0.5.0.post1) with config: model='/home/jie/.cache/modelscope/hub/ZhipuAI/glm-4-9b-chat', speculative_config=None, tokenizer='/home/jie/.cache/modelscope/hub/ZhipuAI/glm-4-9b-chat', skip_tokenizer_init=False, tokenizer_mode=auto, revision=None, rope_scaling=None, rope_theta=None, tokenizer_revision=None, trust_remote_code=True, dtype=torch.bfloat16, max_seq_len=131072, download_dir=None, load_format=LoadFormat.AUTO, tensor_parallel_size=1, disable_custom_all_reduce=False, quantization=None, enforce_eager=True, kv_cache_dtype=auto, quantization_param_path=None, device_config=cuda, decoding_config=DecodingConfig(guided_decoding_backend='outlines'), seed=0, served_model_name=/home/jie/.cache/modelscope/hub/ZhipuAI/glm-4-9b-chat)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING 06-28 16:15:26 tokenizer.py:126] Using a slow tokenizer. This might cause a significant slowdown. Consider using a fast tokenizer instead.\n",
      "INFO 06-28 16:15:32 model_runner.py:160] Loading model weights took 17.5635 GB\n",
      "INFO 06-28 16:15:52 gpu_executor.py:83] # GPU blocks: 10804, # CPU blocks: 6553\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Processed prompts: 100%|██████████| 10/10 [00:00<00:00, 20.44it/s, est. speed input: 4436.22 toks/s, output: 102.22 toks/s]\n"
     ]
    }
   ],
   "source": [
    "llm = LLM(\n",
    "    model=model_name,\n",
    "    tensor_parallel_size=tp_size,\n",
    "    max_model_len=max_model_len,\n",
    "    trust_remote_code=True,\n",
    "    enforce_eager=True,\n",
    "    # GLM-4-9B-Chat-1M 如果遇见 OOM 现象，建议开启下述参数\n",
    "    # enable_chunked_prefill=True,\n",
    "    # max_num_batched_tokens=8192\n",
    ")\n",
    "stop_token_ids = [151329, 151336, 151338]\n",
    "sampling_params = SamplingParams(\n",
    "    temperature=0.95, max_tokens=1024, stop_token_ids=stop_token_ids)\n",
    "\n",
    "inputs = tokenizer.apply_chat_template(\n",
    "    prompt, tokenize=False, add_generation_prompt=True)\n",
    "\n",
    "outputs = llm.generate(prompts=inputs, sampling_params=sampling_params)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "[\"否\"]\n",
      "\n",
      "[\"否\"]\n",
      "\n",
      "[\"否\"]\n",
      "\n",
      "[\"否\"]\n",
      "\n",
      "[\"否\"]\n",
      "\n",
      "[\"否\"]\n",
      "\n",
      "[\"否\"]\n",
      "\n",
      "[\"否\"]\n",
      "\n",
      "[\"否\"]\n",
      "\n",
      "[\"否\"]\n"
     ]
    }
   ],
   "source": [
    "for i in range(10):\n",
    "    print(outputs[i].outputs[0].text)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llm",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
